"""
仿真信号生成器
根据MaintAGT论文中提到的Sim2Real方法实现
包括谐波、边带调制、轴承故障信号等的生成
"""

import numpy as np
import matplotlib.pyplot as plt
from typing import List, Tuple, Dict, Optional
from dataclasses import dataclass


@dataclass
class SignalParams:
    """信号参数配置"""
    fs: float = 5120.0  # 采样频率 Hz
    n_samples: int = 16384  # 固定采样点数 16*1024
    noise_level: float = 0.0  # 噪声水平
    
    @property
    def duration(self) -> float:
        """根据采样点数和采样频率计算时长"""
        return self.n_samples / self.fs


class SignalGenerator:
    """信号生成器基类"""
    
    def __init__(self, params: SignalParams):
        self.params = params
        self.t = np.linspace(0, params.duration, params.n_samples)
        
    def add_noise(self, signal: np.ndarray) -> np.ndarray:
        """添加高斯白噪声"""
        if self.params.noise_level > 0:
            noise = np.random.normal(0, self.params.noise_level, len(signal))
            return signal + noise
        return signal


class HarmonicGenerator(SignalGenerator):
    """谐波信号生成器
    
    谐波信号是故障诊断中最常见的信号类型，通常由旋转机械的不平衡、
    对中不良等故障产生
    """
    
    def generate(self, 
                 fundamental_freq: float,
                 num_harmonics: int = 5,
                 amplitudes: Optional[List[float]] = None,
                 phases: Optional[List[float]] = None) -> Tuple[np.ndarray, Dict]:
        """
        生成谐波信号
        
        参数:
            fundamental_freq: 基频 (Hz)
            num_harmonics: 谐波数量
            amplitudes: 各谐波幅值列表，如果为None则自动衰减
            phases: 各谐波相位列表 (弧度)
            
        返回:
            signal: 生成的信号
            info: 信号信息字典
        """
        if amplitudes is None:
            # 默认幅值按1/n衰减
            amplitudes = [1.0 / (i + 1) for i in range(num_harmonics)]
        
        if phases is None:
            phases = [0.0] * num_harmonics
            
        signal = np.zeros_like(self.t)
        
        # 生成各次谐波并叠加
        for i in range(num_harmonics):
            freq = fundamental_freq * (i + 1)
            signal += amplitudes[i] * np.sin(2 * np.pi * freq * self.t + phases[i])
        
        signal = self.add_noise(signal)
        
        info = {
            'type': '谐波信号',
            'fundamental_freq': fundamental_freq,
            'num_harmonics': num_harmonics,
            'frequencies': [fundamental_freq * (i + 1) for i in range(num_harmonics)],
            'amplitudes': amplitudes
        }
        
        return signal, info


class ModulationGenerator(SignalGenerator):
    """调制（边带）信号生成器
    
    调制信号常见于齿轮故障、轴承外圈故障等，表现为载波频率
    被调制频率调制，产生边带
    """
    
    def generate_am(self,
                    carrier_freq: float,
                    modulation_freq: float,
                    carrier_amp: float = 1.0,
                    modulation_index: float = 0.5) -> Tuple[np.ndarray, Dict]:
        """
        生成调幅(AM)信号
        
        参数:
            carrier_freq: 载波频率 (Hz)
            modulation_freq: 调制频率 (Hz)
            carrier_amp: 载波幅值
            modulation_index: 调制指数 (0-1)
            
        返回:
            signal: 生成的信号
            info: 信号信息字典
        """
        # AM信号: A(1 + m*cos(2πfm*t)) * cos(2πfc*t)
        modulation = 1 + modulation_index * np.cos(2 * np.pi * modulation_freq * self.t)
        carrier = np.cos(2 * np.pi * carrier_freq * self.t)
        signal = carrier_amp * modulation * carrier
        
        signal = self.add_noise(signal)
        
        info = {
            'type': '调幅信号 (AM)',
            'carrier_freq': carrier_freq,
            'modulation_freq': modulation_freq,
            'sideband_freqs': [
                carrier_freq - modulation_freq,
                carrier_freq + modulation_freq
            ],
            'modulation_index': modulation_index
        }
        
        return signal, info
    
    def generate_fm(self,
                    carrier_freq: float,
                    modulation_freq: float,
                    carrier_amp: float = 1.0,
                    frequency_deviation: float = 10.0) -> Tuple[np.ndarray, Dict]:
        """
        生成调频(FM)信号
        
        参数:
            carrier_freq: 载波频率 (Hz)
            modulation_freq: 调制频率 (Hz)
            carrier_amp: 载波幅值
            frequency_deviation: 频率偏移 (Hz)
            
        返回:
            signal: 生成的信号
            info: 信号信息字典
        """
        # FM信号: A*cos(2πfc*t + β*sin(2πfm*t))
        modulation_index = frequency_deviation / modulation_freq
        phase = 2 * np.pi * carrier_freq * self.t + \
                modulation_index * np.sin(2 * np.pi * modulation_freq * self.t)
        signal = carrier_amp * np.cos(phase)
        
        signal = self.add_noise(signal)
        
        info = {
            'type': '调频信号 (FM)',
            'carrier_freq': carrier_freq,
            'modulation_freq': modulation_freq,
            'frequency_deviation': frequency_deviation,
            'modulation_index': modulation_index
        }
        
        return signal, info


class BearingFaultGenerator(SignalGenerator):
    """轴承故障信号生成器
    
    轴承故障信号包含周期性冲击成分，冲击频率取决于轴承几何参数
    和转速。常见的特征频率包括：
    - BPFO: 外圈故障频率
    - BPFI: 内圈故障频率  
    - BSF: 滚动体故障频率
    - FTF: 保持架故障频率
    """
    
    def generate_outer_race_fault(self,
                                   rotation_freq: float,
                                   num_balls: int = 9,
                                   contact_angle: float = 0.0,
                                   impact_amplitude: float = 2.0,
                                   damping: float = 500.0,
                                   resonance_freq: float = 3000.0) -> Tuple[np.ndarray, Dict]:
        """
        生成轴承外圈故障信号
        
        参数:
            rotation_freq: 旋转频率 (Hz)
            num_balls: 滚动体数量
            contact_angle: 接触角 (度)
            impact_amplitude: 冲击幅值
            damping: 阻尼系数
            resonance_freq: 共振频率 (Hz)
            
        返回:
            signal: 生成的信号
            info: 信号信息字典
        """
        # 计算外圈故障特征频率 BPFO
        # BPFO = (n/2) * fr * (1 - (d/D)*cos(α))
        # 简化计算，假设 d/D ≈ 0.2
        bpfo = num_balls * rotation_freq * (1 - 0.2 * np.cos(np.radians(contact_angle))) / 2
        
        signal = self._generate_impact_signal(
            bpfo, impact_amplitude, damping, resonance_freq
        )
        
        info = {
            'type': '轴承外圈故障信号',
            'rotation_freq': rotation_freq,
            'BPFO': bpfo,
            'num_balls': num_balls,
            'resonance_freq': resonance_freq
        }
        
        return signal, info
    
    def generate_inner_race_fault(self,
                                   rotation_freq: float,
                                   num_balls: int = 9,
                                   contact_angle: float = 0.0,
                                   impact_amplitude: float = 2.0,
                                   damping: float = 500.0,
                                   resonance_freq: float = 3000.0) -> Tuple[np.ndarray, Dict]:
        """
        生成轴承内圈故障信号
        
        参数:
            rotation_freq: 旋转频率 (Hz)
            num_balls: 滚动体数量
            contact_angle: 接触角 (度)
            impact_amplitude: 冲击幅值
            damping: 阻尼系数
            resonance_freq: 共振频率 (Hz)
            
        返回:
            signal: 生成的信号
            info: 信号信息字典
        """
        # 计算内圈故障特征频率 BPFI
        # BPFI = (n/2) * fr * (1 + (d/D)*cos(α))
        bpfi = num_balls * rotation_freq * (1 + 0.2 * np.cos(np.radians(contact_angle))) / 2
        
        signal = self._generate_impact_signal(
            bpfi, impact_amplitude, damping, resonance_freq
        )
        
        info = {
            'type': '轴承内圈故障信号',
            'rotation_freq': rotation_freq,
            'BPFI': bpfi,
            'num_balls': num_balls,
            'resonance_freq': resonance_freq
        }
        
        return signal, info
    
    def _generate_impact_signal(self,
                                fault_freq: float,
                                amplitude: float,
                                damping: float,
                                resonance_freq: float) -> np.ndarray:
        """
        生成周期性冲击信号
        
        每个冲击建模为衰减振荡
        """
        signal = np.zeros_like(self.t)
        impact_period = 1.0 / fault_freq
        
        # 生成冲击位置
        num_impacts = int(self.params.duration * fault_freq)
        impact_times = np.arange(num_impacts) * impact_period
        
        # 为每个冲击位置添加小的随机扰动（模拟滑移）
        impact_times += np.random.normal(0, impact_period * 0.02, num_impacts)
        
        # 生成每个冲击
        for impact_time in impact_times:
            if impact_time < self.params.duration:
                # 冲击响应：衰减正弦波
                mask = self.t >= impact_time
                t_relative = self.t[mask] - impact_time
                
                # 添加随机幅值变化
                amp_variation = amplitude * (0.8 + 0.4 * np.random.random())
                
                impact_response = amp_variation * \
                                np.exp(-damping * t_relative) * \
                                np.sin(2 * np.pi * resonance_freq * t_relative)
                
                signal[mask] += impact_response
        
        signal = self.add_noise(signal)
        
        return signal


class CompositeSignalGenerator(SignalGenerator):
    """复合信号生成器
    
    真实的机械振动信号往往是多种信号成分的叠加
    """
    
    def combine_signals(self,
                       signals: List[np.ndarray],
                       weights: Optional[List[float]] = None) -> np.ndarray:
        """
        叠加多个信号
        
        参数:
            signals: 信号列表
            weights: 权重列表，如果为None则等权重
            
        返回:
            combined_signal: 叠加后的信号
        """
        if weights is None:
            weights = [1.0] * len(signals)
        
        combined_signal = np.zeros_like(signals[0])
        
        for signal, weight in zip(signals, weights):
            combined_signal += weight * signal
        
        return combined_signal
    
    def generate_complex_fault_signal(self,
                                      rotation_freq: float = 30.0) -> Tuple[np.ndarray, Dict]:
        """
        生成复杂故障信号示例
        
        包含：不平衡谐波 + 齿轮啮合调制 + 轴承故障冲击
        
        参数:
            rotation_freq: 旋转频率 (Hz)
            
        返回:
            signal: 生成的复合信号
            info: 信号信息字典
        """
        # 1. 不平衡引起的谐波
        harmonic_gen = HarmonicGenerator(self.params)
        harmonic_signal, harmonic_info = harmonic_gen.generate(
            fundamental_freq=rotation_freq,
            num_harmonics=3,
            amplitudes=[1.0, 0.3, 0.1]
        )
        
        # 2. 齿轮啮合调制（假设齿轮有20个齿）
        gear_teeth = 20
        mesh_freq = rotation_freq * gear_teeth
        mod_gen = ModulationGenerator(self.params)
        modulation_signal, mod_info = mod_gen.generate_am(
            carrier_freq=mesh_freq,
            modulation_freq=rotation_freq,
            carrier_amp=0.5,
            modulation_index=0.6
        )
        
        # 3. 轴承外圈故障
        bearing_gen = BearingFaultGenerator(self.params)
        bearing_signal, bearing_info = bearing_gen.generate_outer_race_fault(
            rotation_freq=rotation_freq,
            num_balls=9,
            impact_amplitude=1.5
        )
        
        # 叠加信号
        composite_signal = self.combine_signals(
            [harmonic_signal, modulation_signal, bearing_signal],
            weights=[1.0, 0.8, 1.2]
        )
        
        info = {
            'type': '复合故障信号',
            'components': {
                'harmonic': harmonic_info,
                'modulation': mod_info,
                'bearing_fault': bearing_info
            }
        }
        
        return composite_signal, info


def plot_signal_analysis(t: np.ndarray,
                        signal: np.ndarray,
                        fs: float,
                        title: str = "信号分析"):
    """
    绘制信号的时域和频域图
    
    参数:
        t: 时间数组
        signal: 信号数组
        fs: 采样频率
        title: 图标题
    """
    fig, axes = plt.subplots(2, 1, figsize=(12, 8))
    
    # 时域图
    axes[0].plot(t[:2000], signal[:2000], 'b-', linewidth=0.8)
    axes[0].set_xlabel('时间 (秒)', fontsize=12)
    axes[0].set_ylabel('幅值', fontsize=12)
    axes[0].set_title(f'{title} - 时域波形', fontsize=14, fontweight='bold')
    axes[0].grid(True, alpha=0.3)
    
    # 频域图 (FFT)
    n = len(signal)
    freq = np.fft.fftfreq(n, 1/fs)
    fft_vals = np.fft.fft(signal)
    fft_mag = np.abs(fft_vals) / n
    
    # 只显示正频率部分
    positive_freq_idx = freq > 0
    freq_positive = freq[positive_freq_idx]
    fft_mag_positive = fft_mag[positive_freq_idx]
    
    # 限制频率范围到有意义的部分
    max_display_freq = min(fs/2, 5000)
    display_idx = freq_positive < max_display_freq
    
    axes[1].plot(freq_positive[display_idx], fft_mag_positive[display_idx], 'r-', linewidth=0.8)
    axes[1].set_xlabel('频率 (Hz)', fontsize=12)
    axes[1].set_ylabel('幅值', fontsize=12)
    axes[1].set_title(f'{title} - 频谱', fontsize=14, fontweight='bold')
    axes[1].grid(True, alpha=0.3)
    axes[1].set_xlim([0, max_display_freq])
    
    plt.tight_layout()
    
    return fig


if __name__ == "__main__":
    print("信号生成器模块已加载")
    print("使用示例请参考 demo.py")


